{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# IPEX-LLM on Intel CPU\n",
    "\n",
    "> [IPEX-LLM](https://github.com/intel-analytics/ipex-llm/) is a PyTorch library for running LLM on Intel CPU and GPU (e.g., local PC with iGPU, discrete GPU such as Arc, Flex and Max) with very low latency.\n",
    "\n",
    "This example goes over how to use LlamaIndex to interact with [`ipex-llm`](https://github.com/intel-analytics/ipex-llm/) for text generation and chat on CPU. \n",
    "\n",
    "> **Note**\n",
    ">\n",
    "> You could refer to [here](https://github.com/run-llama/llama_index/tree/main/llama-index-integrations/llms/llama-index-llms-ipex-llm/examples) for full examples of `IpexLLM`. Please note that for running on Intel CPU, please specify `-d 'cpu'` in command argument when running the examples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install `llama-index-llms-ipex-llm`. This will also install `ipex-llm` and its dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-ipex-llm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example we'll use [HuggingFaceH4/zephyr-7b-alpha](https://huggingface.co/HuggingFaceH4/zephyr-7b-alpha) model for demostration. It requires updating `transformers` and `tokenizers` packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -U transformers==4.37.0 tokenizers==0.15.2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before loading the Zephyr model, you'll need to define `completion_to_prompt` and `messages_to_prompt` for formatting prompts. This is essential for preparing inputs that the model can interpret accurately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Transform a string into input zephyr-specific input\n",
    "def completion_to_prompt(completion):\n",
    "    return f\"<|system|>\\n</s>\\n<|user|>\\n{completion}</s>\\n<|assistant|>\\n\"\n",
    "\n",
    "\n",
    "# Transform a list of chat messages into zephyr-specific input\n",
    "def messages_to_prompt(messages):\n",
    "    prompt = \"\"\n",
    "    for message in messages:\n",
    "        if message.role == \"system\":\n",
    "            prompt += f\"<|system|>\\n{message.content}</s>\\n\"\n",
    "        elif message.role == \"user\":\n",
    "            prompt += f\"<|user|>\\n{message.content}</s>\\n\"\n",
    "        elif message.role == \"assistant\":\n",
    "            prompt += f\"<|assistant|>\\n{message.content}</s>\\n\"\n",
    "\n",
    "    # ensure we start with a system prompt, insert blank if needed\n",
    "    if not prompt.startswith(\"<|system|>\\n\"):\n",
    "        prompt = \"<|system|>\\n</s>\\n\" + prompt\n",
    "\n",
    "    # add final assistant prompt\n",
    "    prompt = prompt + \"<|assistant|>\\n\"\n",
    "\n",
    "    return prompt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic Usage\n",
    "\n",
    "Load the Zephyr model locally using IpexLLM using `IpexLLM.from_model_id`. It will load the model directly in its Huggingface format and convert it automatically to low-bit format for inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3665a56541d541a7a1f285953d2c8a2a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/8 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-04-11 21:36:54,739 - INFO - Converting the current model to sym_int4 format......\n"
     ]
    }
   ],
   "source": [
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\n",
    "    \"ignore\", category=UserWarning, message=\".*padding_mask.*\"\n",
    ")\n",
    "\n",
    "from llama_index.llms.ipex_llm import IpexLLM\n",
    "\n",
    "llm = IpexLLM.from_model_id(\n",
    "    model_name=\"HuggingFaceH4/zephyr-7b-alpha\",\n",
    "    tokenizer_name=\"HuggingFaceH4/zephyr-7b-alpha\",\n",
    "    context_window=512,\n",
    "    max_new_tokens=128,\n",
    "    generate_kwargs={\"do_sample\": False},\n",
    "    completion_to_prompt=completion_to_prompt,\n",
    "    messages_to_prompt=messages_to_prompt,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you can proceed to use the loaded model for text completion and interactive chat. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Text Completion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "in a far-off land, \n",
      "there was a young girl named Lily. \n",
      "Lily lived in a small village surrounded by lush green forests and rolling hills. She loved nothing more than spending her days exploring the woods and playing with her animal friends. \n",
      "One day, while wandering through the forest, Lily stumbled upon a magical tree. The tree was unlike any other she had ever seen. Its trunk was made of shimmering crystal, and its branches were adorned with sparkling jewels. \n",
      "Lily was immediately drawn to the tree and sat down to admire its beauty. Suddenly,\n"
     ]
    }
   ],
   "source": [
    "completion_response = llm.complete(\"Once upon a time, \")\n",
    "print(completion_response.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Streaming Text Completion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "who loved to play with her toys. She had a favorite teddy bear named Ted, and a doll named Dolly. She would spend hours playing with them, imagining all sorts of adventures. One day, she decided to take Ted and Dolly on a real adventure. She packed a backpack with some snacks, a blanket, and a map. They set off on a hike in the nearby woods. The little girl was so excited that she could barely contain her joy. Ted and Dolly were happy to be along for the ride. They walked for what seemed like hours, but the little girl didn't mind"
     ]
    }
   ],
   "source": [
    "response_iter = llm.stream_complete(\"Once upon a time, there's a little girl\")\n",
    "for response in response_iter:\n",
    "    print(response.delta, end=\"\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Chat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant: The Big Bang Theory is a popular American sitcom that aired from 2007 to 2019. The show follows the lives of two brilliant but socially awkward physicists, Leonard Hofstadter (Johnny Galecki) and Sheldon Cooper (Jim Parsons), and their friends and colleagues, Penny (Kaley Cuoco), Rajesh Koothrappali (Kunal Nayyar), and Howard Wolowitz (Simon Helberg). The show is set in Pasadena, California, and revolves around the characters' work at Caltech and\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.llms import ChatMessage\n",
    "\n",
    "message = ChatMessage(role=\"user\", content=\"Explain Big Bang Theory briefly\")\n",
    "resp = llm.chat([message])\n",
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Streaming Chat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AI stands for Artificial Intelligence. It refers to the development of computer systems that can perform tasks that typically require human intelligence, such as learning, reasoning, and problem-solving. AI involves the use of machine learning algorithms, natural language processing, and other advanced techniques to enable computers to understand and respond to human input in a more natural and intuitive way."
     ]
    }
   ],
   "source": [
    "message = ChatMessage(role=\"user\", content=\"What is AI?\")\n",
    "resp = llm.stream_chat([message], max_tokens=256)\n",
    "for r in resp:\n",
    "    print(r.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save/Load Low-bit Model\n",
    "Alternatively, you might save the low-bit model to disk once and use `from_model_id_low_bit` instead of `from_model_id` to reload it for later use - even across different machines. It is space-efficient, as the low-bit model demands significantly less disk space than the original model. And `from_model_id_low_bit` is also more efficient than `from_model_id` in terms of speed and memory usage, as it skips the model conversion step."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To save the low-bit model, use `save_low_bit` as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "saved_lowbit_model_path = (\n",
    "    \"./zephyr-7b-alpha-low-bit\"  # path to save low-bit model\n",
    ")\n",
    "\n",
    "llm._model.save_low_bit(saved_lowbit_model_path)\n",
    "del llm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the model from saved lowbit model path as follows. \n",
    "> Note that the saved path for the low-bit model only includes the model itself but not the tokenizers. If you wish to have everything in one place, you will need to manually download or copy the tokenizer files from the original model's directory to the location where the low-bit model is saved."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-04-11 21:38:06,151 - INFO - Converting the current model to sym_int4 format......\n"
     ]
    }
   ],
   "source": [
    "llm_lowbit = IpexLLM.from_model_id_low_bit(\n",
    "    model_name=saved_lowbit_model_path,\n",
    "    tokenizer_name=\"HuggingFaceH4/zephyr-7b-alpha\",\n",
    "    # tokenizer_name=saved_lowbit_model_path,  # copy the tokenizers to saved path if you want to use it this way\n",
    "    context_window=512,\n",
    "    max_new_tokens=64,\n",
    "    completion_to_prompt=completion_to_prompt,\n",
    "    generate_kwargs={\"do_sample\": False},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try stream completion using the loaded low-bit model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A large language model (LLM) is a type of artificial intelligence (AI) model that is trained on a massive amount of text data. These models are capable of generating human-like responses to text inputs and can be used for various natural language processing (NLP) tasks, such as text classification, sentiment analysis"
     ]
    }
   ],
   "source": [
    "response_iter = llm_lowbit.stream_complete(\"What is Large Language Model?\")\n",
    "for response in response_iter:\n",
    "    print(response.delta, end=\"\", flush=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "shane-llamaindex",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
